from enum import Enum
from typing import List, Dict, Tuple, Optional, Any, Union, get_type_hints
import re

class TypeSchema:
    """
    A class representing the schema of data types.
    This is the Python equivalent of the Rust TypeSchema enum.
    """
    class Kind(Enum):
        VOID = "Void"
        INT = "Int"
        FLOAT = "Float"
        STR = "String"
        BOOL = "Bool"
        JSON = "Json"
        LIST = "List"
        MAP = "Map"
        TUPLE = "Tuple" 
        OPT = "Opt"
        DATATYPE = "DataType"
    
    def __init__(self, kind, inner=None, key=None, value=None, elements=None, name=None):
        self.kind = kind
        self.inner = inner  # For List and Opt
        self.key = key      # For Map
        self.value = value  # For Map
        self.elements = elements or []  # For Tuple
        self.name = name    # For DataType
    
    @classmethod
    def Void(cls):
        return cls(cls.Kind.VOID)
    
    @classmethod
    def Int(cls):
        return cls(cls.Kind.INT)
    
    @classmethod
    def Float(cls):
        return cls(cls.Kind.FLOAT)
    
    @classmethod
    def Str(cls):
        return cls(cls.Kind.STR)
    
    @classmethod
    def Bool(cls):
        return cls(cls.Kind.BOOL)
    
    @classmethod
    def Json(cls):
        return cls(cls.Kind.JSON)
    
    @classmethod
    def List(cls, inner_type):
        return cls(cls.Kind.LIST, inner=inner_type)
    
    @classmethod
    def Map(cls, key_type, value_type):
        return cls(cls.Kind.MAP, key=key_type, value=value_type)
    
    @classmethod
    def Tuple(cls, element_types):
        return cls(cls.Kind.TUPLE, elements=element_types)
    
    @classmethod
    def Opt(cls, inner_type):
        return cls(cls.Kind.OPT, inner=inner_type)
    
    @classmethod
    def DataType(cls, name):
        return cls(cls.Kind.DATATYPE, name=name)
    
    def short_str(self):
        """Return a short string representation of the type."""
        if self.kind == self.Kind.VOID:
            return "Void"
        elif self.kind == self.Kind.INT:
            return "Int"
        elif self.kind == self.Kind.FLOAT:
            return "Float"
        elif self.kind == self.Kind.STR:
            return "String"
        elif self.kind == self.Kind.BOOL:
            return "Bool"
        elif self.kind == self.Kind.JSON:
            return "Json"
        elif self.kind == self.Kind.LIST:
            return f"[{self.inner.short_str()}]"
        elif self.kind == self.Kind.MAP:
            return f"Map<{self.key.short_str()}, {self.value.short_str()}>"
        elif self.kind == self.Kind.TUPLE:
            element_strs = [element.short_str() for element in self.elements]
            return f"({', '.join(element_strs)})"
        elif self.kind == self.Kind.OPT:
            return f"?{self.inner.short_str()}"
        elif self.kind == self.Kind.DATATYPE:
            return self.name
    
    def __str__(self):
        """Return a string representation of the type."""
        if self.kind == self.Kind.VOID:
            return "Void"
        elif self.kind == self.Kind.INT:
            return "Int"
        elif self.kind == self.Kind.FLOAT:
            return "Float"
        elif self.kind == self.Kind.STR:
            return "String"
        elif self.kind == self.Kind.BOOL:
            return "Bool"
        elif self.kind == self.Kind.JSON:
            return "Json"
        elif self.kind == self.Kind.LIST:
            return f"List<{self.inner}>"
        elif self.kind == self.Kind.MAP:
            return f"Map<{self.key}, {self.value}>"
        elif self.kind == self.Kind.TUPLE:
            element_strs = [str(element) for element in self.elements]
            return f"({', '.join(element_strs)})"
        elif self.kind == self.Kind.OPT:
            return f"?{self.inner}"
        elif self.kind == self.Kind.DATATYPE:
            return self.name
    
    def __eq__(self, other):
        """Check if two TypeSchema objects are equal."""
        if not isinstance(other, TypeSchema):
            return False
        
        if self.kind != other.kind:
            return False
        
        if self.kind == self.Kind.LIST or self.kind == self.Kind.OPT:
            return self.inner == other.inner
        elif self.kind == self.Kind.MAP:
            return self.key == other.key and self.value == other.value
        elif self.kind == self.Kind.TUPLE:
            if len(self.elements) != len(other.elements):
                return False
            return all(a == b for a, b in zip(self.elements, other.elements))
        elif self.kind == self.Kind.DATATYPE:
            return self.name == other.name
        
        return True

